"""
系统测试脚本
"""

import cv2
import numpy as np
import logging
from main import QueueLengthPredictionSystem
from utils import setup_logging, PerformanceMonitor

def test_with_video(video_path, output_path=None):
    """使用视频测试系统"""
    setup_logging()
    logger = logging.getLogger(__name__)
    
    # 初始化系统
    system = QueueLengthPredictionSystem()
    monitor = PerformanceMonitor()
    
    # 创建视频写入器
    if output_path:
        # 获取第一帧来确定视频尺寸
        cap = cv2.VideoCapture(video_path)
        ret, frame = cap.read()
        if ret:
            out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'XVID'), 25, 
                                (frame.shape[1], frame.shape[0]))
        cap.release()
    else:
        out = None
    
    try:
        # 处理视频
        cap = cv2.VideoCapture(video_path)
        frame_count = 0
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
                
            # 处理帧
            import time
            start_time = time.time()
            
            result = system.process_frame(frame, frame_count)
            
            processing_time = time.time() - start_time
            monitor.record_processing_time(processing_time)
            
            # 可视化
            vis_frame = system.draw_visualization(frame, result)
            
            # 写入输出视频
            if out and vis_frame is not None:
                out.write(vis_frame)
            
            # 显示进度
            if frame_count % 100 == 0:
                logger.info(f"处理帧数: {frame_count}")
                monitor.print_performance_report()
            
            frame_count += 1
            
        cap.release()
        if out:
            out.release()
            
        # 最终性能报告
        logger.info("测试完成")
        monitor.print_performance_report()
        
    except Exception as e:
        logger.error(f"测试失败: {str(e)}")

def test_with_images(image_folder, output_folder=None):
    """使用图像序列测试系统"""
    import os
    import glob
    
    setup_logging()
    logger = logging.getLogger(__name__)
    
    # 初始化系统
    system = QueueLengthPredictionSystem()
    
    # 获取图像文件
    image_files = sorted(glob.glob(os.path.join(image_folder, "*.jpg")) +
                        glob.glob(os.path.join(image_folder, "*.png")))
    
    if not image_files:
        logger.error("未找到图像文件")
        return
    
    results = []
    
    for i, image_file in enumerate(image_files):
        try:
            # 读取图像
            frame = cv2.imread(image_file)
            if frame is None:
                continue
                
            # 处理帧
            result = system.process_frame(frame, i)
            results.append(result)
            
            # 保存结果图像
            if output_folder:
                vis_frame = system.draw_visualization(frame, result)
                output_file = os.path.join(output_folder, f"result_{i:06d}.jpg")
                cv2.imwrite(output_file, vis_frame)
            
            if i % 50 == 0:
                logger.info(f"处理图像: {i}/{len(image_files)}")
                
        except Exception as e:
            logger.error(f"处理图像 {image_file} 失败: {str(e)}")
    
    # 保存结果
    from utils import save_results_to_json
    save_results_to_json(results, "test_results.json")
    
    # 计算性能指标
    from utils import calculate_performance_metrics
    metrics = calculate_performance_metrics(results)
    logger.info(f"测试指标: {metrics}")

if __name__ == "__main__":
    # 测试视频文件
    test_with_video("test_video.mp4", "output_video.avi")
    
    # 或者测试图像序列
    # test_with_images("test_images", "output_images")